package com.example.provider.kafka;

import com.example.api.common.trace.TraceIdUtil;
import com.example.api.common.message.KafkaMessage;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.producer.ProducerInterceptor;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.util.Map;

@Slf4j
public class KafkaTraceIdInterceptor implements ProducerInterceptor<String, String> {
    private static final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public ProducerRecord<String, String> onSend(ProducerRecord<String, String> record) {
        try {
            String value = record.value();
            KafkaMessage<?> message = objectMapper.readValue(value, KafkaMessage.class);
            message.setTraceId(TraceIdUtil.getTraceId());
            String newValue = objectMapper.writeValueAsString(message);
            log.info("Kafka producer traceId: {}", message.getTraceId());
            return new ProducerRecord<>(record.topic(), record.partition(), record.key(), newValue);
        } catch (Exception e) {
            log.error("Failed to add traceId to Kafka message", e);
            return record;
        }
    }

    @Override
    public void onAcknowledgement(RecordMetadata metadata, Exception exception) {
    }

    @Override
    public void close() {
    }

    @Override
    public void configure(Map<String, ?> configs) {
    }
} 